Upload app.py

app.py

@@ -0,0 +1,125 @@
+import gradio as gr
+import cv2
+import numpy as np
+
+class ImageProcessor:
+    @staticmethod
+    def process_image(image):
+        if image is None:
+            print("错误：输入图像为空。")
+            return None, 0
+
+        try:
+            # 确保图像格式正确
+            if len(image.shape) == 2:  # 如果是灰度图
+                image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+            elif image.shape[2] == 4:  # 如果是RGBA
+                image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGR)
+
+            # 边缘保留滤波EPF 去噪
+            blur = cv2.pyrMeanShiftFiltering(image, sp=21, sr=55)
+
+            # 转成灰度图像
+            gray = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)
+
+            # 得到二值图像区间阈值
+            ret, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
+
+            # 距离变换
+            dist = cv2.distanceTransform(binary, cv2.DIST_L2, 3)
+            dist_output = cv2.normalize(dist, None, 0, 1.0, cv2.NORM_MINMAX)
+            ret, surface = cv2.threshold(dist_output, 0.5*dist_output.max(), 255, cv2.THRESH_BINARY)
+
+            # 标记连通区域
+            ret, markers = cv2.connectedComponents(np.uint8(surface))
+            markers = markers + 1
+
+            # 未知区域标记
+            kernel = np.ones((3, 3), np.uint8)
+            unknown = cv2.subtract(cv2.dilate(binary, kernel, iterations=1), np.uint8(surface))
+            markers[unknown == 255] = 0
+
+            # 分水岭算法分割
+            markers = cv2.watershed(image, markers=markers)
+            markers_8u = np.uint8(markers)
+
+            colors = [(255,0,0), (0,255,0), (0,0,255), (255,255,0),
+                      (255,0,255), (0,255,255), (255,128,0), (255,0,128),
+                      (128,255,0), (128,0,255), (255,128,128), (128,255,255)]
+
+            areas = []
+            for i in range(2, np.max(markers) + 1):
+                mask = cv2.inRange(markers_8u, i, i)
+                contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                if contours:
+                    areas.append(cv2.contourArea(contours[0]))
+
+            if not areas:
+                print("警告：未检测到任何对象。")
+                return image, 0
+
+            hist, bin_edges = np.histogram(areas, bins=20)
+            most_common_bin = np.argmax(hist)
+            standard_area = (bin_edges[most_common_bin] + bin_edges[most_common_bin + 1]) / 2
+            area_threshold_low = standard_area * 0.7
+            area_threshold_high = standard_area * 1.3
+
+            object_count = 0
+            for i in range(2, np.max(markers) + 1):
+                mask = cv2.inRange(markers_8u, i, i)
+                contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+                if contours:
+                    area = cv2.contourArea(contours[0])
+                    if area_threshold_low <= area <= area_threshold_high:
+                        object_count += 1
+                    elif area > area_threshold_high:
+                        num_objects = round(area / standard_area)
+                        object_count += num_objects
+                    
+                    color = colors[(i-2)%len(colors)]
+                    cv2.drawContours(image, contours, -1, color, -1)
+                    
+                    M = cv2.moments(contours[0])
+                    if M['m00'] != 0:
+                        cx = int(M['m10']/M['m00'])
+                        cy = int(M['m01']/M['m00'])
+                        cv2.drawMarker(image, (cx,cy), (0,0,255), cv2.MARKER_CROSS, 10, 2)
+
+            cv2.putText(image, f"数量={object_count}", (20,50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0,255,0), 3)
+
+            return image, object_count
+
+        except Exception as e:
+            print(f"图像处理过程中发生错误: {e}")
+            return None, 0
+
+def process_and_count(input_image):
+    if input_image is None:
+        return None, "未上传图像"
+    
+    # 转换图像格式
+    input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
+    
+    processed_image, count = ImageProcessor.process_image(input_image)
+    
+    if processed_image is None:
+        return None, "图像处理错误"
+    
+    # 转换回RGB格式以供Gradio显示
+    processed_image = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
+    
+    return processed_image, f"检测到的对象数量: {count}"
+
+iface = gr.Interface(
+    fn=process_and_count,
+    inputs=gr.Image(),
+    outputs=[
+        gr.Image(label="处理后的图像"),
+        gr.Textbox(label="数量")
+    ],
+    title="螺帽管计数器",
+    description="上传一张图像或拍照以统计螺帽管个数。程序将处理图像并返回检测到的螺帽管数量。"
+)
+
+if __name__ == "__main__":
+    iface.launch(share=True)