Spaces:
Sleeping
Sleeping
| import cv2 | |
| import os.path | |
| # from frame_do_video import images_to_video | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from file_utils import project_dir | |
| def train(): | |
| """ | |
| Funksiya modelni train qiladi | |
| data uyidagi formatda bo'lish kerak | |
| # - splitted | |
| # - train | |
| # - good | |
| # - problem | |
| # - val | |
| # - good | |
| # - problem | |
| """ | |
| data_joyi = 'traffic_laws/scripts/splitted/' | |
| model = YOLO('yolov8n-cls.pt') | |
| model.train(data=data_joyi, epochs=100, imgsz=224, batch=512, save_period=10, device='cuda:0', augment=True) | |
| metrics = model.val() | |
| print(metrics.top1) # top1 aniqligi | |
| def tekshirish(path2,which_mdel=True): | |
| """ | |
| test qilish, model va rasmni berishimiz kerak | |
| """ | |
| train_qilingan_model_joyi = os.path.join( | |
| project_dir(), | |
| "models", | |
| "classification", | |
| "tl-14", | |
| "weights/best.pt" | |
| ) | |
| my_model="tl-14" | |
| if not which_mdel: | |
| my_model="128_1_90" | |
| test_rasm_joyi =(path2) | |
| model_custom = YOLO("models/classification/"+my_model+"/weights/best.pt") | |
| natijalar = model_custom(test_rasm_joyi) # predict on an image | |
| natija = natijalar[0].names[np.argmax(natijalar[0].probs.cpu().numpy().data)] | |
| return (f"Label natija: {natija}") | |
| def process(video_path,which_model): | |
| # saqlash_path = video_path.split('/')[-1].split(".")[0] | |
| saqlash_path = "assets/images" | |
| problem_frame = 0 | |
| good_frame = 0 | |
| print("currnet path: ", os.getcwd()) | |
| if not os.path.exists(saqlash_path): | |
| # Create a new directory because it does not exist | |
| os.makedirs(saqlash_path) | |
| if not os.path.exists(saqlash_path): | |
| raise Exception("Sorry, no numbers below zero") | |
| cap = cv2.VideoCapture(video_path) | |
| saqlash_path="assets/images" | |
| fourcc = cv2.VideoWriter_fourcc(*'XVID') | |
| # out = cv2.VideoWriter('output.avi', fourcc, 20.0, (640,480)) | |
| out = cv2.VideoWriter( | |
| saqlash_path, fourcc, 20.0, (int(cap.get(3)), int(cap.get(4))) | |
| ) | |
| while(cap.isOpened()): | |
| ret, frame = cap.read() | |
| if ret==True: | |
| #frame = cv2.flip(frame,1) | |
| # print(tekshirish(frame)) | |
| if tekshirish(frame,which_model) == "Label natija: good": | |
| font = cv2.FONT_HERSHEY_COMPLEX | |
| cv2.putText(frame, 'good', (0, 100), font, 2, (255, 255, 255), 3) | |
| good_frame += 1 | |
| elif tekshirish(frame, which_model) == "Label natija: problem": | |
| font = cv2.FONT_HERSHEY_COMPLEX | |
| cv2.putText(frame, 'problem', (0, 100), font, 2, (255, 255, 255), 3) | |
| # out.write(frame) | |
| cv2.imwrite(saqlash_path + "/%#05d.jpg" % problem_frame, frame) | |
| problem_frame += 1 | |
| # cv2.imshow('frame' ,frame) | |
| if cv2.waitKey(1) & 0xFF == ord('q'): | |
| break | |
| else: | |
| break | |
| cap.release() | |
| out.release() | |
| cv2.destroyAllWindows() | |
| # images_to_video(saqlash_path, video_name = saqlash_path+'_problem.mp4', fps = 24) | |
| return problem_frame,good_frame,saqlash_path | |