import gdown

def download_file_from_google_drive(file_id, output_file):
    """
    Download a file from Google Drive.

    :param file_id: The Google Drive file ID.
    :param output_file: The name of the file to save.
    """
    url = f"https://drive.google.com/uc?id={file_id}"
    gdown.download(url, output_file, quiet=False)

# Example usage:
file_id = "1Wgh9dWT6SbakJhvuNkSaIa1ydFtkfUW6"
out = "average_model.pth"
download_file_from_google_drive(file_id,out)

from super_gradients.training import models
import torch
import supervision as sv
import gradio as gr 

DEVICE = 'cuda' if torch.cuda.is_available() else "cpu"
MODEL_ARCH = 'yolo_nas_l'
clasess = ["Airplane"]
checkpoint_path= "average_model.pth"


def run(image , CONFIDENCE_TRESHOLD) : 
  best_model = models.get(
      MODEL_ARCH,
      num_classes=len(clasess),
      checkpoint_path= checkpoint_path
  ).to(DEVICE)
  result = list(best_model.predict(image, conf=CONFIDENCE_TRESHOLD))[0]
  detections = sv.Detections(
          xyxy=result.prediction.bboxes_xyxy,
          confidence=result.prediction.confidence,
          class_id=result.prediction.labels.astype(int)
      )
  box_annotator = sv.BoxAnnotator()
  annotated_frame = box_annotator.annotate(
      scene=image.copy(),
      detections=detections,
      labels=clasess
  )
  return annotated_frame
iface = gr.Interface(
    fn=run,
    inputs=[gr.Image(label="Input image", type="numpy") , gr.Slider(0, 1, value=0.5, label="Select your CONFIDENCE_TRESHOLD")],
    outputs=gr.Image(label="The Prediction Output :", type="numpy"),
    title="Aerial Airport YOLO Nas object detection",
    allow_flagging=False , 
    description="I conducted fine-tuning on the YOLO-NAS (YOLO Neural Architecture Search) model, a cutting-edge object detection architecture developed by Deci-AI. My objective was to enhance its ability to detect airplanes in the 'Aerial Airport' dataset",
)
iface.launch(debug=True)