<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
"""
This script is to extract sift features from object data
distributed by Robot Learning Lab, UC Berkeley.
Object data is available from here::

    * http://rll.berkeley.edu/amazon_picking_challenge/


Usage
-----
1. Download dataset to data dir, and extract it.
2. Execute following::

    $ roscore
    $ rosrun imagesift imagesift
    $ rosrun extract_sift_from_objdata.py _object:=oreo_mega_stuf


Attention
---------
You should change dirname for following items manually::

    * kygen_squeakin_eggs_plush_puppies  -> kyjen_squeakin_eggs_plush_puppies
    * rollodex_mesh_collection_jumbo_pencil_cup -> rolodex_jumbo_pencil_cup

"""
import os
import sys
import collections
import gzip
import cPickle as pickle

import cv2
import numpy as np
import yaml
import progressbar

import rospy
from sensor_msgs.msg import CameraInfo
from posedetection_msgs.srv import Feature0DDetect
from posedetection_msgs.msg import ImageFeature0D

from sift_matcher_oneimg import SiftMatcherOneImage


def get_train_imgpaths(obj_name):
    """Find train image paths from data/obj_name"""
    dirname = os.path.dirname(os.path.abspath(__file__))
    obj_dir = os.path.join(dirname, '../data/', obj_name)
    if not os.path.exists(obj_dir):
        rospy.logwarn('Object data does not exists: {}'.format(obj_name))
        return
    os.chdir(obj_dir)
    imgpaths = []
    for imgfile in os.listdir('.'):
        if not imgfile.endswith('.jpg'):
            continue
        raw_path = os.path.join(obj_dir, imgfile)
        maskfile = os.path.splitext(imgfile)[0] + '_mask.pbm'
        mask_path = os.path.join(obj_dir, 'masks', maskfile)
        imgpaths.append((raw_path, mask_path))
    os.chdir(dirname)
    return imgpaths


def save_siftdata(siftdata, obj_name):
    """Save sift data to data/siftdata/{obj_name}.pkl.gz"""
    dirname = os.path.dirname(os.path.abspath(__file__))
    siftdata_dir = os.path.join(dirname, '../data/siftdata')
    if not os.path.exists(siftdata_dir):
        os.mkdir(siftdata_dir)
    filename = os.path.join(siftdata_dir, obj_name+'.pkl.gz')
    with gzip.open(filename, 'wb') as f:
        pickle.dump(siftdata, f)


def extract_sift_from_objdata(obj_name):
    """Extract sift data from object images"""
    positions = []
    descriptors = []
    imgpaths = get_train_imgpaths(obj_name)
    if imgpaths is None:
        return   # skip if img does not exists
    progress = progressbar.ProgressBar(widgets=['{o}: '.format(o=obj_name),
        progressbar.Bar(), progressbar.Percentage(), ' ', progressbar.ETA()])
    for raw_path, mask_path in progress(imgpaths):
        raw_img = cv2.imread(raw_path)
        mask_img = cv2.imread(mask_path)
        train_img = cv2.add(mask_img, raw_img)
        train_features = SiftMatcherOneImage.imgsift_client(train_img)
        train_pos = np.array(train_features.positions)
        train_des = np.array(train_features.descriptors)
        positions.append(train_pos)
        descriptors.append(train_des)
    positions, descriptors = map(np.array, [positions, descriptors])
    siftdata = dict(positions=positions, descriptors=descriptors)
    # save sift data
    save_siftdata(siftdata, obj_name)


def main():
    rospy.init_node('extract_sift_from_objdata')
    pub = rospy.Publisher('/camera_info', CameraInfo, queue_size=1)
    pub.publish()  # to enable imagesift service

    dirname = os.path.dirname(os.path.abspath(__file__))
    ymlfile = os.path.join(dirname, '../data/object_list.yml')
    all_objects = yaml.load(open(ymlfile))

    obj_names = rospy.get_param('~object',
                                'oreo_mega_stuf,safety_works_safety_glasses')
    obj_names = obj_names.split(',')
    if len(obj_names) == 1 and obj_names[0] == 'all':
        obj_names = all_objects
    rospy.loginfo('objects: {obj}'.format(obj=obj_names))<commit_msg>match changed class name<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
"""
This script is to extract sift features from object data
distributed by Robot Learning Lab, UC Berkeley.
Object data is available from here::

    * http://rll.berkeley.edu/amazon_picking_challenge/


Usage
-----
1. Download dataset to data dir, and extract it.
2. Execute following::

    $ roscore
    $ rosrun imagesift imagesift
    $ rosrun extract_sift_from_objdata.py _object:=oreo_mega_stuf


Attention
---------
You should change dirname for following items manually::

    * kygen_squeakin_eggs_plush_puppies  -> kyjen_squeakin_eggs_plush_puppies
    * rollodex_mesh_collection_jumbo_pencil_cup -> rolodex_jumbo_pencil_cup

"""
import os
import sys
import collections
import gzip
import cPickle as pickle

import cv2
import numpy as np
import yaml
import progressbar

import rospy
from sensor_msgs.msg import CameraInfo
from posedetection_msgs.srv import Feature0DDetect
from posedetection_msgs.msg import ImageFeature0D

from sift_matcher_oneimg import SiftMatcherOneImg


def get_train_imgpaths(obj_name):
    """Find train image paths from data/obj_name"""
    dirname = os.path.dirname(os.path.abspath(__file__))
    obj_dir = os.path.join(dirname, '../data/', obj_name)
    if not os.path.exists(obj_dir):
        rospy.logwarn('Object data does not exists: {}'.format(obj_name))
        return
    os.chdir(obj_dir)
    imgpaths = []
    for imgfile in os.listdir('.'):
        if not imgfile.endswith('.jpg'):
            continue
        raw_path = os.path.join(obj_dir, imgfile)
        maskfile = os.path.splitext(imgfile)[0] + '_mask.pbm'
        mask_path = os.path.join(obj_dir, 'masks', maskfile)
        imgpaths.append((raw_path, mask_path))
    os.chdir(dirname)
    return imgpaths


def save_siftdata(siftdata, obj_name):
    """Save sift data to data/siftdata/{obj_name}.pkl.gz"""
    dirname = os.path.dirname(os.path.abspath(__file__))
    siftdata_dir = os.path.join(dirname, '../data/siftdata')
    if not os.path.exists(siftdata_dir):
        os.mkdir(siftdata_dir)
    filename = os.path.join(siftdata_dir, obj_name+'.pkl.gz')
    with gzip.open(filename, 'wb') as f:
        pickle.dump(siftdata, f)


def extract_sift_from_objdata(obj_name):
    """Extract sift data from object images"""
    positions = []
    descriptors = []
    imgpaths = get_train_imgpaths(obj_name)
    if imgpaths is None:
        return   # skip if img does not exists
    progress = progressbar.ProgressBar(widgets=['{o}: '.format(o=obj_name),
        progressbar.Bar(), progressbar.Percentage(), ' ', progressbar.ETA()])
    for raw_path, mask_path in progress(imgpaths):
        raw_img = cv2.imread(raw_path)
        mask_img = cv2.imread(mask_path)
        train_img = cv2.add(mask_img, raw_img)
        train_features = SiftMatcherOneImg.imgsift_client(train_img)
        train_pos = np.array(train_features.positions)
        train_des = np.array(train_features.descriptors)
        positions.append(train_pos)
        descriptors.append(train_des)
    positions, descriptors = map(np.array, [positions, descriptors])
    siftdata = dict(positions=positions, descriptors=descriptors)
    # save sift data
    save_siftdata(siftdata, obj_name)


def main():
    rospy.init_node('extract_sift_from_objdata')
    pub = rospy.Publisher('/camera_info', CameraInfo, queue_size=1)
    pub.publish()  # to enable imagesift service

    dirname = os.path.dirname(os.path.abspath(__file__))
    ymlfile = os.path.join(dirname, '../data/object_list.yml')
    all_objects = yaml.load(open(ymlfile))

    obj_names = rospy.get_param('~object',
                                'oreo_mega_stuf,safety_works_safety_glasses')
    obj_names = obj_names.split(',')
    if len(obj_names) == 1 and obj_names[0] == 'all':
        obj_names = all_objects
    rospy.loginfo('objects: {obj}'.format(obj=obj_names))
{'fraction_non_alphanumeric': 0.11343513409470037, 'fraction_numerical': 0.0033363274733735405, 'mean_word_length': 3.573943661971831, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1462684', 'n_tokens_mistral': 2716, 'n_tokens_neox': 2623, 'n_words': 548}