main.py

import json
import os
import subprocess
import time

import cv2
import requests
import torch
from flask import Flask, Response, request
from ultralytics import YOLO
#from predict import predictImg
from flask_socketio import SocketIO, emit

# 导入钢材缺陷检测预测器
from predict.steelDefectPredictor import SteelDefectPredictor
from predict.yolo_engine import YoloEngine
# from predict.efficientvit_sam import Efficientvit_Sam
# 添加PIL导入，用于图像处理
from PIL import Image
import numpy as np
import cv2


# Flask 应用设置
class VideoProcessingApp:
    def __init__(self, host='0.0.0.0', port=5000):
        """初始化 Flask 应用并设置路由"""
        self.app = Flask(__name__)
        self.socketio = SocketIO(self.app, cors_allowed_origins="*")  # 初始化 SocketIO
        self.host = host
        self.port = port
        self.setup_routes()
        # self.Efficientvit_Sam = Efficientvit_Sam()
        self.YoloEngine = YoloEngine
        self.data = {}  # 存储接收参数
        self.paths = {
            'download': './runs/video/download.mp4',
            'output': './runs/video/output.mp4',
            'camera_output': "./runs/video/camera_output.avi",
            'video_output': "./runs/video/camera_output.avi"
        }
        self.recording = False  # 标志位，判断是否正在录制视频

    def setup_routes(self):
        """设置所有路由"""
        self.app.add_url_rule('/file_names', 'file_names', self.file_names, methods=['GET'])
        self.app.add_url_rule('/predictImg', 'predictImg', self.predictImg, methods=['POST'])
        self.app.add_url_rule('/predictVideo', 'predictVideo', self.predictVideo)
        self.app.add_url_rule('/predictCamera', 'predictCamera', self.predictCamera)
        self.app.add_url_rule('/stopCamera', 'stopCamera', self.stopCamera, methods=['GET'])

        # 添加 WebSocket 事件
        @self.socketio.on('connect')
        def handle_connect():
            print("WebSocket connected!")
            emit('message', {'data': 'Connected to WebSocket server!'})

        @self.socketio.on('disconnect')
        def handle_disconnect():
            print("WebSocket disconnected!")

    def run(self):
        """启动 Flask 应用"""
        self.socketio.run(self.app, host=self.host, port=self.port, allow_unsafe_werkzeug=True)

    def file_names(self):
        """模型列表接口"""
        try:
            with open('./weight_config.json', 'r', encoding='utf-8') as f:
                weight_config = json.load(f)
            
            weight_items = []
            for name in self.get_file_names("./weights"):
                if name in weight_config:
                    weight_items.append({
                        'value': name,
                        'label': weight_config[name]['display_name'],
                        'description': weight_config[name]['description']
                    })
                else:
                    weight_items.append({'value': name, 'label': name})
            
            return json.dumps({'weight_items': weight_items}, ensure_ascii=False)
        except Exception as e:
            print(f"读取权重配置文件失败: {e}")
            # 如果配置文件读取失败，回退到原来的方式
            weight_items = [{'value': name, 'label': name} for name in self.get_file_names("./weights")]
            return json.dumps({'weight_items': weight_items})

    def predictImg(self):
        """图片预测接口"""
        data = request.get_json()
        print(data)
        self.data.clear()
        self.data.update({
            "username": data['username'], "weight": data['weight'],
            "startTime": data['startTime'],
            "inputImg": data['inputImg'],
            "kind": data['kind']
        })
        print(self.data)
        # 下载图片到本地临时文件
        temp_image_path = "./runs/temp_input.jpg"
        os.makedirs("./runs", exist_ok=True)  # 确保目录存在
        self.download(self.data["inputImg"], temp_image_path)

        # 根据权重文件选择检测方法
        if self.data["weight"] == "yolov8.pt":

            # 使用YOLO模型进行检测
            model = YOLO(f'./weights/{self.data["weight"]}', task='detect').to('cuda')
            print(model.device)
            # 读取图片并放大
            img = cv2.imread(temp_image_path)
            scale = 10
            img = cv2.resize(img, dsize=None, fx=scale, fy=scale, interpolation=cv2.INTER_LINEAR)
            start_time = time.time()
            # 执行检测
            yolo_results = model(temp_image_path, conf=0.3)
            result = yolo_results[0]

            # 获取检测框和类别信息
            boxes = result.boxes.xyxy.cpu().numpy()
            cls = result.boxes.cls.cpu().numpy()
            names = result.names

            # 绘制检测结果
            for box, cls_id in zip(boxes, cls):
                x1, y1, x2, y2 = map(int, box * scale)
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                label = names[int(cls_id)]
                cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

            # 保存结果图片
            cv2.imwrite('./runs/result.jpg', img)
            ellipsis_time = time.time() - start_time
            # 构造返回结果
            results = {
                'labels': list(set([names[int(c)] for c in cls])),
                'confidences': [f"{100:.2f}%" for _ in range(len(set([names[int(c)] for c in cls])))],
                'allTime': f"{ellipsis_time}秒"
            }
        elif self.data["weight"] == "Yolov8.engine" or self.data["weight"] == "NEUyolov8n.engine":
            results = self.YoloEngine().predict(temp_image_path)
        elif self.data["weight"] == "Efficientvit_Sam.pth":
            results = self.Efficientvit_Sam.predict(temp_image_path)
        else:
            # 使用钢材缺陷检测预测器
            predict = SteelDefectPredictor(
                weights_path=f'./weights/{self.data["weight"]}',
                img_path=temp_image_path,  # 使用本地临时文件路径
                save_path='./runs/result.jpg'
            )

            # 执行预测
            results = predict.predict()
        uploadedUrl = self.upload('./runs/result.jpg')

        # 删除临时文件
        if os.path.exists(temp_image_path):
            os.remove(temp_image_path)

        if results['labels'] != '预测失败':
            self.data["status"] = 200
            self.data["message"] = "预测成功"
            self.data["outImg"] = uploadedUrl
            self.data["allTime"] = results['allTime']
            self.data["confidence"] = json.dumps(results['confidences'])
            self.data["label"] = json.dumps(results['labels'])
        else:
            self.data["status"] = 400
            self.data["message"] = "该图片无法识别，请重新上传！"
        path = self.data["inputImg"].split('/')[-1]
        # if os.path.exists('./' + path):
        # os.remove('./' + path)
        return json.dumps(self.data, ensure_ascii=False)

    def predictVideo(self):
        """视频流处理接口"""
        self.data.clear()
        self.data.update({
            "username": request.args.get('username'), "weight": request.args.get('weight'),
            "conf": request.args.get('conf'), "startTime": request.args.get('startTime'),
            "inputVideo": request.args.get('inputVideo'),
            "kind": request.args.get('kind')
        })
        self.download(self.data["inputVideo"], self.paths['download'])
        cap = cv2.VideoCapture(self.paths['download'])
        if not cap.isOpened():
            raise ValueError("无法打开视频文件")
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        print(fps)

        # 视频写入器
        video_writer = cv2.VideoWriter(
            self.paths['video_output'],
            cv2.VideoWriter_fourcc(*'XVID'),
            fps,
            (640, 480)
        )

        # 使用DeeplabV3模型
        # sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'fuchuang'))
        from deeplab import DeeplabV3

        model = DeeplabV3(
            model_path=f'./weights/{self.data["weight"]}',
            num_classes=4,  # 背景 + 3种缺陷
            mix_type=0
        )

        def generate():
            try:
                while cap.isOpened():
                    ret, frame = cap.read()
                    if not ret:
                        break
                    frame = cv2.resize(frame, (640, 480))
                    # 转换为PIL图像
                    pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                    # 使用DeeplabV3进行预测
                    result_image = model.detect_image(pil_image)
                    # 转换回OpenCV格式
                    processed_frame = cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR)
                    video_writer.write(processed_frame)
                    _, jpeg = cv2.imencode('.jpg', processed_frame)
                    yield b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n'
            finally:
                self.cleanup_resources(cap, video_writer)
                self.socketio.emit('message', {'data': '处理完成，正在保存！'})
                for progress in self.convert_avi_to_mp4(self.paths['video_output']):
                    self.socketio.emit('progress', {'data': progress})
                uploadedUrl = self.upload(self.paths['output'])
                self.data["outVideo"] = uploadedUrl
                self.save_data(json.dumps(self.data), 'http://localhost:9999/videoRecords')
                self.cleanup_files([self.paths['download'], self.paths['output'], self.paths['video_output']])

        return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')

    def predictCamera(self):
        """摄像头视频流处理接口"""
        self.data.clear()
        self.data.update({
            "username": request.args.get('username'), "weight": request.args.get('weight'),
            "kind": request.args.get('kind'),
            "conf": request.args.get('conf'), "startTime": request.args.get('startTime')
        })
        self.socketio.emit('message', {'data': '正在加载，请稍等！'})

        # 使用DeeplabV3模型
        # sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'fuchuang'))
        from deeplab import DeeplabV3

        model = DeeplabV3(
            model_path=f'./weights/{self.data["weight"]}',
            num_classes=4,  # 背景 + 3种缺陷
            mix_type=0
        )

        cap = cv2.VideoCapture(0)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        video_writer = cv2.VideoWriter(self.paths['camera_output'], cv2.VideoWriter_fourcc(*'XVID'), 20, (640, 480))
        self.recording = True

        def generate():
            try:
                while self.recording:
                    ret, frame = cap.read()
                    if not ret:
                        break
                    # 转换为PIL图像
                    pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                    # 使用DeeplabV3进行预测
                    result_image = model.detect_image(pil_image)
                    # 转换回OpenCV格式
                    processed_frame = cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR)
                    if self.recording and video_writer:
                        video_writer.write(processed_frame)
                    _, jpeg = cv2.imencode('.jpg', processed_frame)
                    yield b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + jpeg.tobytes() + b'\r\n'
            finally:
                self.cleanup_resources(cap, video_writer)
                self.socketio.emit('message', {'data': '处理完成，正在保存！'})
                for progress in self.convert_avi_to_mp4(self.paths['camera_output']):
                    self.socketio.emit('progress', {'data': progress})
                uploadedUrl = self.upload(self.paths['output'])
                self.data["outVideo"] = uploadedUrl
                print(self.data)
                self.save_data(json.dumps(self.data), 'http://localhost:9999/cameraRecords')
                self.cleanup_files([self.paths['download'], self.paths['output'], self.paths['camera_output']])

        return Response(generate(), mimetype='multipart/x-mixed-replace; boundary=frame')

    def stopCamera(self):
        """停止摄像头预测"""
        self.recording = False
        return json.dumps({"status": 200, "message": "预测成功", "code": 0})

    def save_data(self, data, path):
        """将结果数据上传到服务器"""
        headers = {'Content-Type': 'application/json'}
        try:
            response = requests.post(path, data=data, headers=headers)
            print("记录上传成功！" if response.status_code == 200 else f"记录上传失败，状态码: {response.status_code}")
        except requests.RequestException as e:
            print(f"上传记录时发生错误: {str(e)}")

    def convert_avi_to_mp4(self, temp_output):
        """使用 FFmpeg 将 AVI 格式转换为 MP4 格式，并显示转换进度。"""
        ffmpeg_command = f"ffmpeg -i {temp_output} -vcodec libx264 {self.paths['output']} -y"
        process = subprocess.Popen(ffmpeg_command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                   text=True)
        total_duration = self.get_video_duration(temp_output)

        for line in process.stderr:
            if "time=" in line:
                try:
                    time_str = line.split("time=")[1].split(" ")[0]
                    h, m, s = map(float, time_str.split(":"))
                    processed_time = h * 3600 + m * 60 + s
                    if total_duration > 0:
                        progress = (processed_time / total_duration) * 100
                        yield progress
                except Exception as e:
                    print(f"解析进度时发生错误: {e}")

        process.wait()
        yield 100

    def get_video_duration(self, path):
        """获取视频总时长（秒）"""
        try:
            cap = cv2.VideoCapture(path)
            if not cap.isOpened():
                return 0
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            fps = cap.get(cv2.CAP_PROP_FPS)
            cap.release()
            return total_frames / fps if fps > 0 else 0
        except Exception:
            return 0

    def get_file_names(self, directory):
        """获取指定文件夹中的所有文件名"""
        try:
            return [file for file in os.listdir(directory) if os.path.isfile(os.path.join(directory, file))]
        except Exception as e:
            print(f"发生错误: {e}")
            return []

    def upload(self, out_path):
        """上传处理后的图片或视频文件到远程服务器"""
        upload_url = "http://localhost:9999/files/upload"
        try:
            with open(out_path, 'rb') as file:
                files = {'file': (os.path.basename(out_path), file)}
                response = requests.post(upload_url, files=files)
                if response.status_code == 200:
                    print("文件上传成功！")
                    return response.json()['data']
                else:
                    print("文件上传失败！")
        except Exception as e:
            print(f"上传文件时发生错误: {str(e)}")

    def download(self, url, save_path):
        """下载文件并保存到指定路径"""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        try:
            with requests.get(url, stream=True) as response:
                response.raise_for_status()
                with open(save_path, 'wb') as file:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            file.write(chunk)
            print(f"文件已成功下载并保存到 {save_path}")
        except requests.RequestException as e:
            print(f"下载失败: {e}")

    def cleanup_files(self, file_paths):
        """清理文件"""
        for path in file_paths:
            if os.path.exists(path):
                os.remove(path)

    def cleanup_resources(self, cap, video_writer):
        """释放资源"""
        if cap.isOpened():
            cap.release()
        if video_writer is not None:
            video_writer.release()
        cv2.destroyAllWindows()


# 启动应用
if __name__ == '__main__':
    video_app = VideoProcessingApp()
    video_app.run()