Update app.py

app.py

@@ -201,34 +201,37 @@
 # iface.launch()
 
 import gradio as gr
-import cv2
-from ultralytics import YOLO
-
-model = YOLO('Model_IV.pt')
-print(model.names)
-# webcamera = cv2.VideoCapture(0)
-# webcamera.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
-# webcamera.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
-
-# while True:
-#     success, frame = webcamera.read()
+import torch
+from PIL import Image
+import torchvision.transforms as T
+
+# Load your model
+model = torch.load("model.pt")
+model.eval()
+
+# Define preprocessing
+transform = T.Compose([
+    T.Resize((224, 224)),  # Adjust to your model's input size
+    T.ToTensor(),
+])
+
+def predict(image):
+    # Preprocess the image
+    img_tensor = transform(image).unsqueeze(0)  # Add batch dimension
     
-#     results = model.track(frame, conf=0.2, imgsz=480)
-#     cv2.putText(frame, f"Total: {len(results[0].boxes)}", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
-#     cv2.imshow("Live Camera", results[0].plot())
-
-#     if cv2.waitKey(1) == ord('q'):
-#         break
-
-def run_model(img):
-    results = model(img)
-    annotated_img = results.render()[0]
-    return annotated_img
-
-# webcamera.release()
-# cv2.destroyAllWindows()
-
-
-
-iface = gr.Interface(fn=run_model, inputs="webcam", outputs=annotated_img)
-iface.launch()
+    # Make prediction
+    with torch.no_grad():
+        output = model(img_tensor)
+    
+    # Process output (adjust based on your model's format)
+    return output.tolist()  # or post-process the results as needed
+
+# Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),  # Accepts image input
+    outputs="json"  # Customize based on your output format
+)
+
+if __name__ == "__main__":
+    demo.launch()