Update app.py

app.py

@@ -141,37 +141,41 @@
 # gradio_app.launch(debug=True)
 
 # make sure you have the following dependencies
+import gradio as gr
 import torch
-import numpy as np
-from models.common import DetectMultiBackend
-from utils.general import non_max_suppression, scale_boxes
-from utils.torch_utils import select_device, smart_inference_mode
-from utils.augmentations import letterbox
-import PIL.Image
-from huggingface_hub import hf_hub_download 
-hf_hub_download("SakshiRathi77/void-space-detection", filename="weights/best.pt", local_dir="./")
-
-
-@smart_inference_mode()
-def predict(image_path, weights='best.pt', imgsz=640, conf_thres=0.1, iou_thres=0.45):
-    # Initialize
-    device = select_device('0')
-    # model = DetectMultiBackend(weights='best.pt', device="0", fp16=False, data='data/coco.yaml')
-    model = DetectMultiBackend(weights='best.pt', device="0", fp16=False)
-    stride, names, pt = model.stride, model.names, model.pt
-
-    # Load image
-    image = np.array(PIL.Image.open(image_path))
-    img = letterbox(img0, imgsz, stride=stride, auto=True)[0]
-    img = img[:, :, ::-1].transpose(2, 0, 1)
-    img = np.ascontiguousarray(img)
-    img = torch.from_numpy(img).to(device).float()
-    img /= 255.0
-    if img.ndimension() == 3:
-        img = img.unsqueeze(0)
-
-    # Inference
-    # pred = model(img, augment=False, visualize=False)
-    pred = model(img)
-    # Apply NMS
-    pred = non_max_suppression(pred[0][0], conf_thres, iou_thres, classes=None, max_det=1000)
+from torchvision import transforms
+from PIL import Image
+
+# Load the YOLOv9 model
+model_path = "best.pt"  # Replace with the path to your YOLOv9 model
+model = torch.load(model_path)
+
+# Define preprocessing transforms
+preprocess = transforms.Compose([
+    transforms.Resize((640, 640)),  # Resize image to model input size
+    transforms.ToTensor(),           # Convert image to tensor
+])
+
+# Define a function to perform inference
+def detect_void(image):
+    # Preprocess the input image
+    image = Image.fromarray(image)
+    image = preprocess(image).unsqueeze(0)  # Add batch dimension
+
+    # Perform inference
+    with torch.no_grad():
+        output = model(image)
+
+    # Post-process the output if needed
+    # For example, draw bounding boxes on the image
+
+    # Convert the image back to numpy array
+    # and return the result
+    return output.squeeze().numpy()
+
+# Define Gradio interface components
+input_image = gr.inputs.Image(shape=(640, 640), label="Input Image")
+output_image = gr.outputs.Image(label="Output Image")
+
+# Create Gradio interface
+gr.Interface(fn=detect_void, inputs=input_image, outputs=output_image, title="Void Detection App").launch()