Update app.py

app.py

@@ -2,18 +2,15 @@ import os
 import importlib.util
 from huggingface_hub import hf_hub_download
 import gradio as gr
-import torch
-from PIL import Image
-import torchvision.transforms as transforms
 
 # --- CONFIG ---
 PRIVATE_DATASET_ID = "abdulrafay9/containeralign-private"
-TOKEN = os.environ.get("HF_TOKEN")  # secret key
+TOKEN = os.environ.get("HF_TOKEN")
 
 if not TOKEN:
     raise RuntimeError("HF_TOKEN is not set. Add it in Settings → Variables and secrets → Secrets.")
 
-# --- DOWNLOAD PRIVATE FILES ---
+# --- DOWNLOAD FILES ---
 core_path = hf_hub_download(
     repo_id=PRIVATE_DATASET_ID,
     repo_type="dataset",
@@ -27,7 +24,7 @@ weights_path = hf_hub_download(
     token=TOKEN,
 )
 
-# --- LOAD CORE FILE ---
+# --- LOAD MODULE ---
 spec = importlib.util.spec_from_file_location("app_core", core_path)
 app_core = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(app_core)
@@ -35,28 +32,19 @@ spec.loader.exec_module(app_core)
 # --- LOAD MODEL ---
 model = app_core.load_model(weights_path)
 
-# --- IMAGE TRANSFORM (same as training) ---
-transform = transforms.Compose([
-    transforms.Resize((128, 128)),
-    transforms.ToTensor(),
-    transforms.Normalize((0.5,), (0.5,))
-])
-
-# --- PREDICT FUNCTION ---
+# --- DEFINE PREDICTION FUNCTION ---
 def predict(image):
-    img = Image.open(image).convert("RGB")
-    tensor = transform(img).unsqueeze(0)
-    result = app_core.predict_alignment(model, tensor)
-    return f"Prediction: {result}"
+    return app_core.predict_alignment(model, image)
 
-# --- GRADIO APP ---
+# --- GRADIO INTERFACE ---
 demo = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="filepath", label="Upload an image"),
-    outputs="text",
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=gr.Textbox(label="Prediction Result"),
     title="Container Alignment Detection",
-    description="Upload an image to check whether containers are aligned or not."
+    description="Upload an image to check whether containers are Aligned or Not Aligned."
 )
 
+# --- RUN APP ---
 if __name__ == "__main__":
     demo.launch()