Update app.py

app.py

@@ -17,11 +17,13 @@ results_cache = {}
 
 # Download model from Hugging Face
 def download_model():
+    print("Downloading model...")
     model_path = hf_hub_download(repo_id="jays009/Restnet50", filename="pytorch_model.bin")
     return model_path
 
 # Load the model from Hugging Face
 def load_model(model_path):
+    print("Loading model...")
     model = models.resnet50(pretrained=False)
     model.fc = nn.Linear(model.fc.in_features, num_classes)
     model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
@@ -76,40 +78,24 @@ def predict_from_url(url):
         print(f"Error during URL processing: {e}")
         return {"error": f"Failed to process the URL: {str(e)}"}
 
-# Main prediction function without event ID for direct uploads
-def predict_direct_upload(image):
+# Main prediction function with caching
+def predict(image, url):
     try:
+        print("Starting prediction...")
         if image:
             result = predict_from_image(image)
-            return result
+        elif url:
+            result = predict_from_url(url)
         else:
-            return {"error": "No image provided. Please upload an image."}
-    except Exception as e:
-        print(f"Error in direct upload prediction function: {e}")
-        return {"error": str(e)}
+            result = {"error": "No input provided. Please upload an image or provide a URL."}
 
-# Main prediction function with caching for paths via Postman
-def predict_with_event_id(data):
-    try:
-        path = data[0].get('path', None)
-        if path:
-            if path.startswith("http://") or path.startswith("https://"):
-                result = predict_from_url(path)
-            elif os.path.isfile(path):
-                image = Image.open(path)
-                result = predict_from_image(image)
-            else:
-                result = {"error": "Invalid path format. Please provide a valid URL or local file path."}
-            
-            event_id = id(result)
-            results_cache[event_id] = result
-            
-            print(f"Event ID: {event_id}, Result: {result}")
-            return {"event_id": event_id, "result": result}
-        else:
-            return {"error": "No path provided. Please provide a valid path."}
+        event_id = id(result)  # Use Python's id() function to generate a unique identifier
+        results_cache[event_id] = result
+
+        print(f"Event ID: {event_id}, Result: {result}")
+        return {"event_id": event_id, "result": result}
     except Exception as e:
-        print(f"Error in prediction function with event ID: {e}")
+        print(f"Error in prediction function: {e}")
         return {"error": str(e)}
 
 # Function to retrieve result by event_id
@@ -144,8 +130,8 @@ with iface:
     submit_button = gr.Button("Submit")
     
     submit_button.click(
-        fn=predict_direct_upload,
-        inputs=[image_input],
+        fn=predict,
+        inputs=[image_input, url_input],
         outputs=output
     )