Update app.py

app.py

@@ -1,84 +1,83 @@
-import gradio as gr
-from PIL import Image
-import torch
-import torchvision.transforms as transforms
-from model import RetinaNet  # Import your RetinaNet model definition
-import cv2
-import numpy as np
-
-# Define the image transformation pipeline
-image_transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-])
-
-# Load the model
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model = RetinaNet(num_classes=2).to(device)
-model.load_state_dict(torch.load("retinanet_best_model.pth", map_location=device))
-model.eval()
-
-# Prediction function
-def predict_image(image):
-    # Preprocess the image
-    img = Image.fromarray(image).convert('RGB')  # Convert Gradio input to PIL Image
-    input_tensor = image_transform(img).unsqueeze(0).to(device)
-
-    # Perform inference
-    with torch.no_grad():
-        prediction = model(input_tensor.float())
-        sum_value = abs(torch.sum(prediction[0]))
-        p_true = abs(prediction[0][0])
-        p_false = abs(prediction[0][1])
-
-    # Interpret the prediction
-    if p_true > 0.7:
-        result = "Accepted"
-        confidence = float(p_true)
-    else:
-        result = "Rejected"
-        confidence = float(p_false)
-
-    return f"Result: {result}, Confidence: {confidence:.2f}"
-
-# Prediction function
-def predict_frame(frame):
-    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    resized_frame = cv2.resize(rgb_frame, (224, 224))
-    normalized_frame = resized_frame / 255.0
-    input_frame = np.expand_dims(normalized_frame, axis=0)
-
-    # Convert to PyTorch tensor and move to device
-    input_frame = torch.from_numpy(input_frame).to(device).float()
-
-    # Permute dimensions to [batch_size, channels, height, width]
-    input_frame = input_frame.permute(0, 3, 1, 2)
-
-    # Predict using the best model
-    with torch.no_grad():
-        prediction = model(input_frame)
-        sum_value=torch.sum(abs(prediction[0]))
-        p_true=abs(prediction[0][0])
-        p_false=abs(prediction[0][1])
-
-    if p_true < 0.4:#if p_true > p_false:
-        result = "Accepted"
-        confidence = float(p_true)
-    else:
-        result = "Rejected"
-        confidence = float(p_false)
-
-    return f"Result: {result}, Confidence: {confidence:.2f}"
-
-# Create the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# RetinaNet Model Prediction")
-    with gr.Row():
-        image_input = gr.Image(label="Upload Image", type="numpy")
-        output_text = gr.Textbox(label="Prediction Result")
-    predict_button = gr.Button("Predict")
-    predict_button.click(predict_image, inputs=image_input, outputs=output_text)
-
-# Launch the app
+import gradio as gr
+from PIL import Image
+import torch
+import torchvision.transforms as transforms
+from model import RetinaNet  # Import your RetinaNet model definition
+import cv2
+import numpy as np
+
+# Define the image transformation pipeline
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Load the model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = RetinaNet(num_classes=2).to(device)
+model.load_state_dict(torch.load("retinanet_best_model.pth", map_location=device))
+model.eval()
+
+# Prediction function
+def predict_image(image, isFrame):
+
+    if isFrame == False:
+        # Preprocess the image
+        img = Image.fromarray(image).convert('RGB')  # Convert Gradio input to PIL Image
+        input_tensor = image_transform(img).unsqueeze(0).to(device)
+
+        # Perform inference
+        with torch.no_grad():
+            prediction = model(input_tensor.float())
+            sum_value = abs(torch.sum(prediction[0]))
+            p_true = abs(prediction[0][0])
+            p_false = abs(prediction[0][1])
+
+        # Interpret the prediction
+        if p_true > 0.7:
+            result = "Accepted"
+            confidence = float(p_true)
+        else:
+            result = "Rejected"
+            confidence = float(p_false)
+    else:
+        frame = image
+        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        resized_frame = cv2.resize(rgb_frame, (224, 224))
+        normalized_frame = resized_frame / 255.0
+        input_frame = np.expand_dims(normalized_frame, axis=0)
+
+        # Convert to PyTorch tensor and move to device
+        input_frame = torch.from_numpy(input_frame).to(device).float()
+
+        # Permute dimensions to [batch_size, channels, height, width]
+        input_frame = input_frame.permute(0, 3, 1, 2)
+
+        # Predict using the best model
+        with torch.no_grad():
+            prediction = model(input_frame)
+            sum_value=torch.sum(abs(prediction[0]))
+            p_true=abs(prediction[0][0])
+            p_false=abs(prediction[0][1])
+
+        if p_true < 0.4:#if p_true > p_false:
+            result = "Accepted"
+            confidence = float(p_true)
+        else:
+            result = "Rejected"
+            confidence = float(p_false)
+
+    return f"Result: {result}, Confidence: {confidence:.2f}"
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# RetinaNet Model Prediction")
+    with gr.Row():
+        image_input = gr.Image(label="Upload Image", type="numpy")
+        output_text = gr.Textbox(label="Prediction Result")
+    predict_button = gr.Button("Predict")
+    predict_button.click(predict_image, inputs=image_input, outputs=output_text)
+
+# Launch the app
 demo.launch()