Create use_with_UI.py

use_with_UI.py

@@ -0,0 +1,113 @@
+import streamlit as st
+import torch
+import requests
+from io import BytesIO
+from PIL import Image
+from torchvision import transforms
+from transformers import ResNetForImageClassification
+
+# --- 1. UI Configuration ---
+# 'centered' ensures the app doesn't stretch across massive screens
+st.set_page_config(page_title="GyroScope Rotation Corrector", layout="centered", page_icon="🔄")
+
+# --- 2. Model Caching ---
+# @st.cache_resource prevents reloading the model every time the user interacts with the UI
+@st.cache_resource
+def load_model():
+    model = ResNetForImageClassification.from_pretrained("LH-Tech-AI/GyroScope")
+    model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    return model, device
+
+model, device = load_model()
+
+# --- 3. Preprocessing & Logic ---
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+ANGLES = [0, 90, 180, 270]
+
+def predict_and_correct(img):
+    # Ensure image is RGB
+    img = img.convert("RGB")
+    tensor = preprocess(img).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        logits = model(pixel_values=tensor).logits
+        probs = torch.softmax(logits, dim=1)[0]
+        pred = probs.argmax().item()
+
+    detected = ANGLES[pred]
+    correction = (360 - detected) % 360
+    
+    # Apply correction (PIL rotate is counter-clockwise)
+    corrected_img = img.rotate(correction, expand=True)
+    
+    # Format probabilities for the UI
+    prob_dict = {f"{a}°": f"{p:.4f}" for a, p in zip(ANGLES, probs)}
+    
+    return corrected_img, detected, correction, prob_dict
+
+# --- 4. Frontend Layout ---
+st.title("🔄 Auto Rotation Corrector")
+st.markdown("Upload an image or provide a URL to automatically fix its orientation.")
+
+st.divider()
+
+# Input Selection
+input_method = st.radio("Select Image Source:", ["Upload a File", "Enter Image URL"], horizontal=True)
+
+img = None
+
+# Input Handling
+if input_method == "Upload a File":
+    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        img = Image.open(uploaded_file)
+else:
+    url = st.text_input("Enter Image URL:", placeholder="https://example.com/image.jpg")
+    if url:
+        try:
+            response = requests.get(url, timeout=5)
+            img = Image.open(BytesIO(response.content))
+        except Exception as e:
+            st.error(f"Could not load image from URL. Error: {e}")
+
+# Preview & Processing Section
+if img:
+    st.divider()
+    
+    # Use columns to keep the UI compact and side-by-side
+    col_left, col_right = st.columns(2)
+    
+    with col_left:
+        st.subheader("Input Preview")
+        st.image(img, use_container_width=True)
+        
+        # The primary action button
+        process_btn = st.button("✨ Correct Rotation", type="primary", use_container_width=True)
+        
+    with col_right:
+        st.subheader("Output Preview")
+        
+        if process_btn:
+            with st.spinner("Analyzing..."):
+                corrected_img, detected, correction, prob_dict = predict_and_correct(img)
+                
+                # Show result
+                st.image(corrected_img, use_container_width=True)
+                
+                # Show stats
+                st.success(f"✅ Detected: **{detected}°** | Correction: **{correction}°**")
+                
+                # Hidden expander for clean UI, but available if the user wants details
+                with st.expander("📊 View Probability Details"):
+                    st.json(prob_dict)
+        else:
+            # Placeholder container before the button is clicked
+            st.info("Waiting for processing... Click the button on the left to correct the rotation.")