Merge pull request SkalskiP#284 from SkalskiP/fix/281_support_for_yol…

…ov5_p6_models fix/281_support_for_yolov5_p6_models
PowZone · Oct 13, 2022 · d0f8db4 · d0f8db4
2 parents b48ce89 + 8e8bab4
commit d0f8db4
Showing 1 changed file with 29 additions and 1 deletion.
diff --git a/src/ai/YOLOV5ObjectDetector.ts b/src/ai/YOLOV5ObjectDetector.ts
@@ -9,12 +9,16 @@ import {Notification} from '../data/enums/Notification';
 import {submitNewNotification} from '../store/notifications/actionCreators';
 import {LabelsSelector} from '../store/selectors/LabelsSelector';
 import {AIYOLOObjectDetectionActions} from '../logic/actions/AIYOLOObjectDetectionActions';
+import {ImageData} from '../store/labels/types';
+import {ImageRepository} from '../logic/imageRepository/ImageRepository';
 
 export class YOLOV5ObjectDetector {
  private static model: YOLOv5;
 
  public static loadModel(modelConfig: ModelConfig, onSuccess?: () => any, onFailure?: () => any) {
- load(modelConfig)
+ const activeImageData: ImageData = LabelsSelector.getActiveImageData();
+ const image = ImageRepository.getById(activeImageData.id)
+ YOLOV5ObjectDetector.loadModelSafely(modelConfig, image)
  .then((model: YOLOv5) => {
  YOLOV5ObjectDetector.model = model;
  store.dispatch(updateYOLOV5ObjectDetectorStatus(true));
@@ -26,10 +30,32 @@ export class YOLOV5ObjectDetector {
  if (onSuccess) onSuccess()
  })
  .catch((error) => {
+ // tslint:disable-next-line:no-console
+ console.log(error)
  if (onFailure) onFailure()
  })
  }
 
+ private static loadModelSafely(modelConfig: ModelConfig, image: HTMLImageElement): Promise<YOLOv5> {
+ return new Promise( (resolve, reject) => {
+ load(modelConfig, [640, 640])
+ .then((model640: YOLOv5) => {
+ model640.detect(image)
+ .then((detections: DetectedObject[]) => resolve(model640))
+ .catch((error: Error) => {
+ load(modelConfig, [1280, 1280])
+ .then((model1280: YOLOv5) => {
+ model1280.detect(image)
+ .then((detections: DetectedObject[]) => resolve(model1280))
+ .catch(reject)
+ })
+ .catch(reject)
+ })
+ })
+ .catch(reject)
+ });
+ }
+
  public static predict(image: HTMLImageElement, callback?: (predictions: DetectedObject[]) => any) {
  if (!YOLOV5ObjectDetector.model) return;
 
@@ -41,6 +67,8 @@ export class YOLOV5ObjectDetector {
  }
  })
  .catch((error) => {
+ // tslint:disable-next-line:no-console
+ console.log(error)
  // TODO: Introduce central logging system like Sentry
  store.dispatch(
  submitNewNotification(