Update app.py

app.py

@@ -16,62 +16,52 @@ if 'models_loaded' not in st.session_state:
     st.session_state.models_loaded = True
 
 # Define image processing and classification functions
-def topwear(encoding,top_wear_model):
-    # Make prediction
+def topwear(encoding):
     outputs = st.session_state.top_wear_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    # Print the result
     return st.session_state.top_wear_model.config.id2label[predicted_class_idx]
 
-def patterns(encoding,pattern_model):
-    # Make prediction
+def patterns(encoding):
     outputs = st.session_state.pattern_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    # Print the result
     return st.session_state.pattern_model.config.id2label[predicted_class_idx]
 
-def prints(encoding,print_model):
-    # Make prediction
+def prints(encoding):
     outputs = st.session_state.print_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    # Print the result
     return st.session_state.print_model.config.id2label[predicted_class_idx]
 
-def sleevelengths(encoding,sleeve_length_model):
-    # Make prediction
+def sleevelengths(encoding):
     outputs = st.session_state.sleeve_length_model(**encoding)
     logits = outputs.logits
     predicted_class_idx = logits.argmax(-1).item()
-    # Print the result
     return st.session_state.sleeve_length_model.config.id2label[predicted_class_idx]
 
-def imageprocessing(url):
-    response = requests.get(url)
-    if response.status_code == 200:
-        image = Image.open(BytesIO(response.content))
-        encoding = image_processor(image.convert("RGB"), return_tensors="pt")
+def imageprocessing(image):
+    encoding = st.session_state.image_processor(image.convert("RGB"), return_tensors="pt")
     return encoding
-    
+
 # Define the function that will be used in each thread
-def call_model(func, encoding, model, results, index):
-    results[index] = func(encoding, model)
+def call_model(func, encoding, results, index):
+    results[index] = func(encoding)
+
 # Run all models in parallel
-def pipes(imagepath):
+def pipes(image):
     # Process the image once and reuse the encoding
-    encoding = imageprocessing(imagepath)
+    encoding = imageprocessing(image)
 
     # Prepare a list to store the results from each thread
     results = [None] * 4
 
     # Create threads for each function call
     threads = [
-        threading.Thread(target=call_model, args=(topwear, encoding, top_wear_model, results, 0)),
-        threading.Thread(target=call_model, args=(patterns, encoding, pattern_model, results, 1)),
-        threading.Thread(target=call_model, args=(prints, encoding, print_model, results, 2)),
-        threading.Thread(target=call_model, args=(sleevelengths, encoding, sleeve_length_model, results, 3)),
+        threading.Thread(target=call_model, args=(topwear, encoding, results, 0)),
+        threading.Thread(target=call_model, args=(patterns, encoding, results, 1)),
+        threading.Thread(target=call_model, args=(prints, encoding, results, 2)),
+        threading.Thread(target=call_model, args=(sleevelengths, encoding, results, 3)),
     ]
 
     # Start all threads
@@ -97,20 +87,19 @@ st.title("Clothing Classification Pipeline")
 
 url = st.text_input("Paste image URL here...")
 if url:
-    response = requests.get(url)
-    if response.status_code == 200:
-        image = Image.open(BytesIO(response.content))
-        encoding = image_processor(image.convert("RGB"), return_tensors="pt")
-        st.image(image.resize((200, 200)), caption="Uploaded Image", use_column_width=False)
-        
-        start_time = time.time()
-
-        try:
-            result = pipes(url)
+    try:
+        response = requests.get(url)
+        if response.status_code == 200:
+            image = Image.open(BytesIO(response.content))
+            st.image(image.resize((200, 200)), caption="Uploaded Image", use_column_width=False)
+            
+            start_time = time.time()
+
+            result = pipes(image)
             st.write("Classification Results (JSON):")
             st.json(result)  # Display results in JSON format
             st.write(f"Time taken: {time.time() - start_time:.2f} seconds")
-        except Exception as e:
-            st.error(f"Error processing the image: {str(e)}")
-    else:
-        st.error("Failed to load image from URL. Please check the URL.")
+        else:
+            st.error("Failed to load image from URL. Please check the URL.")
+    except Exception as e:
+        st.error(f"Error processing the image: {str(e)}")