Beispiel #1
0
    def select_input(self):
        """
        Save the name of the input file, clear output label, set progress bar to zero and actualize input label.
        When Select Input button is trigerred the method is invoked.
        """
        full_paths = QFileDialog.getOpenFileNames(
            None, 'Open File', '',
            '*.jpg *.jpeg *.png *.bmp *.mp4 *.avi *.wmv *.mov *.mkv')
        if full_paths:
            # Clear input files list
            self.files.clear()
            self.output_name = None
            full_path = full_paths[0]
            print(full_paths)
            file = File(self.tmp_dir, full_path[0])

            # Clear output label image, progress bar and label text
            self.output_label.clear()
            self.progress_bar.setValue(0)
            self.path_label.setText(file.get_output_name(self.output_type))

            # Actualize input label image
            if file.type == InputType.IMAGE:
                image = file.full_path
            elif file.type == InputType.VIDEO:
                ret, image = cv2.VideoCapture(file.full_path).read()
                if not ret:
                    raise ValueError
                image = npimg_to_pixmap(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            else:
                raise ValueError
            self.actualize_input_label(image)

            # Add input file to a list of files
            self.files.append(file)
    def process_image(self, window, file, output_type):
        """
        Process input image.
        :param window: PyQt window
        :param file: the name of the input file
        :param output_type: the type of the output
        :return: the name of the output file
        """
        image = read_image_bgr(file.full_path)
        objects_detected = self.detect_objects(image)
        output_name = file.get_output_name(output_type)

        if output_type == OutputType.BORDERS:
            output_image = self.draw_bounding_boxes(image, objects_detected)

        elif output_type == OutputType.PANORAMA:
            objects_cropped = self.crop_objects(image, objects_detected)
            matcher = Matcher(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            matcher.process_objects(0, objects_cropped)
            matcher.save_objects(output_name.split('.')[0] + '-dir/', 0, objects_cropped)
            output_image = matcher.get_panorama(InputType.IMAGE)
        else:
            raise ValueError

        # Save output image to tmp/
        image = Image.fromarray(output_image)
        image.save(output_name)

        # Visualize output image
        window.actualize_output_label(npimg_to_pixmap(output_image))

        return output_name
Beispiel #3
0
    def run_recognition(self):
        """
        Run recognition, process input file.
        When Run Recognition button is trigerred the method is invoked.
        """
        if len(self.files) == 0:
            self.show_message_box(
                MessageBoxType.ERROR, 'Error',
                'There is no input image or video to be processed.')
        elif self.output_type is None:
            self.show_message_box(MessageBoxType.ERROR, 'Error',
                                  'Output type has to be specified.')
        else:
            self.output_label.clear()
            self.output_label.repaint()

            for file in self.files[::-1]:
                # Reset progress bar and output label
                self.progress_bar.setValue(0)
                self.app.processEvents()
                self.path_label.setText(file.get_output_name(self.output_type))

                # Actualize input label image because of input stacking
                if file.type == InputType.IMAGE:
                    image = file.full_path
                elif file.type == InputType.VIDEO:
                    ret, image = cv2.VideoCapture(file.full_path).read()
                    if not ret:
                        raise ValueError
                    image = npimg_to_pixmap(
                        cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
                else:
                    raise ValueError
                self.actualize_input_label(image)

                if file.type == InputType.IMAGE:
                    self.output_name = self.detector.process_image(
                        self, file, self.output_type)

                elif file.type == InputType.VIDEO:
                    self.output_name = self.detector.process_video(
                        self, file, self.output_type)

                else:
                    raise ValueError

                self.progress_bar.setValue(100)
                self.app.processEvents()

                if file.type == InputType.VIDEO:
                    self.show_message_box(MessageBoxType.INFORMATION,
                                          'Information', 'Detection is done')
    def process_video(self, window, file, output_type):
        """
        Process input video.
        :param window: PyQt window
        :param file: the name of the input file
        :param output_type: the type of the output
        :return: the name of the output file
        """
        # Read input video capture
        in_cap = cv2.VideoCapture(file.full_path)

        # Extract video info (fps, number of frames, resolution)
        fps = int(in_cap.get(cv2.CAP_PROP_FPS))
        frame_count = int(in_cap.get(cv2.CAP_PROP_FRAME_COUNT))
        resolution = (int(in_cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(in_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))

        output_name = file.get_output_name(output_type)

        # If output is video capture of frames with object boundaries
        if output_type == OutputType.BORDERS:
            # Create output video capture for writing frames
            out_cap = cv2.VideoWriter(output_name, cv2.VideoWriter_fourcc(*'XVID'), fps, resolution)

        # If output is a panorama image with trajectories of objects
        elif output_type == OutputType.PANORAMA:
            # First frame of input video is gonna be used for drawing trajectories
            ret, image = cv2.VideoCapture(file.full_path).read()
            if not ret:
                raise ValueError
            matcher = Matcher(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

        else:
            raise ValueError

        # Calculate number of steps for progress bar
        step = 100 / frame_count
        frame_cnt = 0
        completed = 0

        # Create temporary directory for saving processed frames with boundaries
        tmp_frame_dir = file.get_frame_dir(output_type)
        create_folder(tmp_frame_dir)

        # Loop for reading whole video frame by frame
        while in_cap.isOpened():

            # Read single frame
            ret, image = in_cap.read()
            if not ret:
                break

            # Visualize input frame which is gonna be processed
            window.actualize_input_label(npimg_to_pixmap(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)))

            # Detect objects in input frame
            objects_detected = self.detect_objects(image)

            # Draw objects boundaries to the actual frame
            image_with_borders = self.draw_bounding_boxes(image, objects_detected)

            # If output is video with object boundaries
            if output_type == OutputType.BORDERS:
                # Write frame with object boundaries to output video
                out_cap.write(cv2.cvtColor(image_with_borders, cv2.COLOR_BGR2RGB))

            # If output is panorama image
            elif output_type == OutputType.PANORAMA:
                # Crop objects in the frame (get numpy array represantition of objects)
                objects_cropped = self.crop_objects(image, objects_detected)

                # Process objects in the frame
                matcher.process_objects(frame_cnt, objects_cropped)

                # Save objects which were detected in the frame
                matcher.save_objects(output_name.split('.')[0] + '-dir/', frame_cnt, objects_cropped)

            # Save visualized frame to tmp/image-dir/
            image = Image.fromarray(image_with_borders)
            image.save(tmp_frame_dir + file.get_frame_name(output_type) + '-frame' + str(frame_cnt) + '.jpg')

            # Visualize processed frame with detected objects to output label
            window.actualize_output_label(npimg_to_pixmap(image_with_borders))

            # Update progress bar
            completed += step
            window.progress_bar.setValue(completed)

            frame_cnt += 1

        # Release source vid
        in_cap.release()
        if output_type == OutputType.BORDERS:
            out_cap.release()

        # Process panorama image and get its name
        elif output_type == OutputType.PANORAMA:
            img = matcher.get_panorama(InputType.VIDEO)
            window.actualize_output_label(npimg_to_pixmap(img))
            img = Image.fromarray(img)
            img.save(output_name)

        cv2.destroyAllWindows()

        return output_name