Upload app.py

app.py

@@ -1,55 +1,60 @@
-from ultralytics import YOLO
-from PIL import Image
+import numpy as np
+import tensorflow as tf
 import gradio as gr
+from tensorflow.keras.preprocessing import image
 from huggingface_hub import snapshot_download
 import os
 
 
+# Load the model from Hugging Face Hub
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
-    print(download_dir)
-    path = os.path.join(download_dir, "best_int8_openvino_model")
-    print(path)
-    detection_model = YOLO(path, task='detect')
-    return detection_model
+    model_path = os.path.join(download_dir, "full_model.weights.h5")
+    model = tf.keras.models.load_model(model_path)
+    return model
 
 
-def predict(pilimg, confidence, iou):
-    source = pilimg
-    result = detection_model.predict(source, conf=confidence, iou=iou)
-    img_bgr = result[0].plot()
-    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
-    return out_pilimg
+# Function to preprocess the uploaded image
+def preprocess_image(img, target_size=(224, 224)):
+    img = img.resize(target_size)  # Resize to match model input
+    img_array = image.img_to_array(img)
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)
+    return img_array
 
 
-REPO_ID = "skngew/9053220B"
-detection_model = load_model(REPO_ID)
+# Perform inference
+def predict(image_input):
+    class_names = ["Defective Tyre", "Good Tyre"]
+    
+    # Preprocess image
+    img_array = preprocess_image(image_input)
+
+    # Get prediction
+    prediction = model.predict(img_array)[0][0]  # Scalar sigmoid output
+    predicted_class_idx = int(prediction >= 0.5)  # 0 if <0.5, 1 if >=0.5
+    predicted_class = class_names[predicted_class_idx]  # Get class name
+
+    return f"Predicted Class: {predicted_class} (Confidence: {prediction:.5f})"
+
+
+# Hugging Face Model Repository ID
+REPO_ID = "skngew/9053220B"  # Change this to your actual repo ID
+
+# Load the model
+model = load_model(REPO_ID)
 
 # Student ID
 student_id = "Student ID: 9053220B"
 
 # Create the Gradio interface
-def create_interface():
-    # Persistent state for default values
-    confidence_default = gr.State(0.5)
-    iou_default = gr.State(0.6)
-
-    interface = gr.Interface(
-        fn=predict,
-        inputs=[
-            gr.Image(type="pil", label="Input Image"),
-            gr.Slider(0, 1, value=confidence_default.value, label="Confidence Threshold"),  # Default to 0.5
-            gr.Slider(0, 1, value=iou_default.value, label="IOU Threshold")  # Default to 0.6
-        ],
-        outputs=gr.Image(type="pil", label="Output Image"),
-        title="Object Detection with YOLOv8",
-        description=student_id,
-        live=False,
-    )
-
-    return interface
-
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload an Image"),
+    outputs=gr.Textbox(label="Prediction"),
+    title="Binary Classification: Good vs. Defective Tire",
+    description=student_id,
+)
 
 # Launch the Gradio app
-app_interface = create_interface()
-app_interface.launch(share=True)
+interface.launch(share=True)