| from ultralytics import YOLO |
| import cv2 |
| import streamlit as st |
| import os |
| import numpy as np |
|
|
| st.set_page_config(layout="wide") |
| model = None |
| user_inputs = {} |
|
|
| with st.sidebar: |
| st.title("Calculate Product Costs from Images") |
| st.write("Simplify pricing with AI-powered image recognition and price summation.") |
| st.text("1.Upload Model (YOLO or etc.)\n2.Set The price (Default is 0)\n3.Upload An Image\n(Optional) 4.Set The Model Confidence") |
|
|
| model_file = st.file_uploader("Upload your model file (.pt)", type=["pt"]) |
| |
| if model_file is not None: |
| |
| temp_model_path = os.path.join("./", model_file.name) |
| with open(temp_model_path, "wb") as f: |
| f.write(model_file.getbuffer()) |
| st.success(f"Model file '{model_file.name}' uploaded and saved successfully!") |
| |
| try: |
| model = YOLO(temp_model_path) |
| st.success("Model loaded successfully!") |
| class_names = model.names |
|
|
| st.write("Enter The Prices:") |
| for idx, name in class_names.items(): |
| user_input = st.text_input(f"Class {idx}: {name}", key=f"class_{idx}") |
| user_inputs[idx] = user_input |
|
|
| if 'collected_list' not in st.session_state: |
| st.session_state.collected_list = [] |
|
|
| if st.button("Submit"): |
| st.write("Collected Inputs:") |
| st.session_state.collected_list = [] |
| for idx in range(len(class_names)): |
| if user_inputs[idx] == "": |
| user_inputs[idx] = 0 |
| elif not user_inputs[idx].isdigit(): |
| user_inputs[idx] = 0 |
| st.session_state.collected_list.append(int(user_inputs[idx])) |
| st.write(st.session_state.collected_list) |
| except Exception as e: |
| st.error(f"Error loading model: {e}") |
| else: |
| st.warning("Please upload a model file that ends with .pt") |
| |
| if model != None: |
| |
| st.subheader("Image Display") |
| image_placeholder = st.empty() |
|
|
| uploaded_image = st.file_uploader("Upload an image to display", type=["png", "jpg", "jpeg"], key="image") |
| conf_str = st.text_input(f"Model Confidence (Default is 0.5)") |
| if uploaded_image is not None: |
| file_bytes = np.asarray(bytearray(uploaded_image.read()), dtype=np.uint8) |
| img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) |
| |
| if st.button("Predict The Image"): |
| if model is not None and uploaded_image is not None: |
| if st.session_state.collected_list != []: |
| try: |
| conf_ = float(conf_str) |
| except Exception as e: |
| if not isinstance(conf_str, float): |
| conf_ = 0.5 |
| |
| results = model.predict(source=img, conf=conf_) |
| if 'sum_price' not in st.session_state: |
| st.session_state.sum_price = 0 |
| |
| for result in results: |
| for box in result.boxes: |
| |
| x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int) |
| cls = int(box.cls[0]) |
| conf = box.conf[0] |
| st.write(f"{result.names[int(box.cls[0])]} : {st.session_state.collected_list[int(box.cls[0])]}") |
| st.session_state.sum_price += st.session_state.collected_list[int(box.cls[0])] |
| label = f"{result.names[int(box.cls[0])]}: {conf:.2f}" |
| cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) |
| cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) |
| |
| st.subheader(f"Sum Price: {st.session_state.sum_price}") |
| st.image(img, channels="BGR", caption="Uploaded Image") |
| st.session_state.sum_price = 0 |
| else: |
| st.warning("Please Submit The Price") |
| else: |
| st.warning("Please Upload an Image") |
| |
|
|
