Spaces:
Runtime error
Runtime error
| 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) |